import type { Base } from "../Base";
import { mat4 } from "wgpu-matrix";
import { GUI } from "dat.gui";
import { quitIfWebGPUNotAvailable } from "../../utils";
import texturedQuadWGSL from "./texturedQuad.wgsl";
import {
  color,
  alpha,
  constant,
  clear,
  type Preset,
  presets,
} from "./defaultConfigs";
import { set16ToRGBNormalized } from "./utils";

class GUIColorHelper {
  normalizedColor: number[];

  constructor(normalizedColor: number[]) {
    this.normalizedColor = normalizedColor;
  }
  get value() {
    return this.normalizedColor.map((v) => Math.round(v * 255));
  }
  set value(rgb255Color: number[]) {
    this.normalizedColor.forEach(
      (_, i) => (this.normalizedColor[i] = rgb255Color[i] / 255)
    );
  }
}
export default class Blending implements Base {
  name = "Blending";
  description =
    "This example provides shows how to use blending in WebGPU. \
  It draws a texture with diagonal lines on top of a canvas with a CSS based checkerboard background. \
  It then draws a texture with 3 blurry circles on top the first texture with <a href='https://gpuweb.github.io/gpuweb/#color-target-state'>blending settings.</a> \
  This lets you see both the effect of blending settings in WebGPU and the final result when composited on top of the canvas.\
  See this <a href='https://webgpufundamentals.org/webgpu/lessons/webgpu-transparency.html'>article</a> for a more detailed explanation. The presets are equivalent to the 2d canvas context's globalCompositingOperation.";
  gui?: GUI;
  settings: {
    alphaMode: GPUCanvasAlphaMode;
    textureSet: string;
    preset: Preset;
    color: {
      operation: GPUBlendOperation;
      srcFactor: GPUBlendFactor;
      dstFactor: GPUBlendFactor;
    };
    alpha: {
      operation: GPUBlendOperation;
      srcFactor: GPUBlendFactor;
      dstFactor: GPUBlendFactor;
    };
    constant: {
      color: number[];
      alpha: number;
    };
    clear: {
      color: number[];
      alpha: number;
      premultiply: boolean;
    };
  } = {
    alphaMode: "premultiplied",
    textureSet: "premultiplied alpha",
    preset: "premultiplied blend (source-over)",
    color: {
      operation: "add",
      srcFactor: "one",
      dstFactor: "one-minus-src-alpha",
    },
    alpha: {
      operation: "add",
      srcFactor: "one",
      dstFactor: "one-minus-src-alpha",
    },
    constant,
    clear,
  };
  render: (() => void) | null = null;
  applyPreset!: () => void;
  async run(canvas: HTMLCanvasElement) {
    const self = this;
    const adapter = await navigator.gpu?.requestAdapter({
      featureLevel: "compatibility",
    });
    const device = await adapter?.requestDevice();
    if (!adapter || !device) return;
    quitIfWebGPUNotAvailable(adapter, device);

    // creates a CSS hsl string from 3 normalized numbers (0 to 1)
    const hsl = (hue: number, saturation: number, lightness: number) =>
      `hsl(${(hue * 360) | 0}, ${saturation * 100}%, ${(lightness * 100) | 0}%)`;

    // creates a CSS hsla string from 4 normalized numbers (0 to 1)
    const hsla = (
      hue: number,
      saturation: number,
      lightness: number,
      alpha: number
    ) =>
      // prettier-ignore
      `hsla(${(hue * 360) | 0}, ${saturation * 100}%, ${(lightness * 100) | 0}%, ${alpha})`;

    // Generates a canvas with 3 circles of different colors with blurred edges.
    function createSourceImage(size: number) {
      const canvasSrc = document.createElement("canvas");
      canvasSrc.width = size;
      canvasSrc.height = size;
      const ctx = canvasSrc.getContext("2d");
      ctx?.translate(size / 2, size / 2);

      if (ctx) ctx.globalCompositeOperation = "screen";
      const numCircles = 3;
      for (let i = 0; i < numCircles; ++i) {
        ctx?.rotate((Math.PI * 2) / numCircles);
        ctx?.save();
        ctx?.translate(size / 6, 0);
        ctx?.beginPath();

        const radius = size / 3;
        ctx?.arc(0, 0, radius, 0, Math.PI * 2);

        const gradient = ctx?.createRadialGradient(
          0,
          0,
          radius / 2,
          0,
          0,
          radius
        );
        const h = i / numCircles;
        gradient?.addColorStop(0.5, hsla(h, 1, 0.5, 1));
        gradient?.addColorStop(1, hsla(h, 1, 0.5, 0));

        if (ctx && gradient) {
          ctx.fillStyle = gradient;
          ctx.fill();
        }
        ctx?.restore();
      }
      return canvasSrc;
    }

    // Generates a canvas with alternating colored and transparent stripes
    function createDestinationImage(size: number) {
      const canvasDst = document.createElement("canvas");
      canvasDst.width = size;
      canvasDst.height = size;
      const ctx = canvasDst.getContext("2d");

      const gradient = ctx?.createLinearGradient(0, 0, size, size);
      for (let i = 0; i <= 6; ++i) {
        gradient?.addColorStop(i / 6, hsl(i / -6, 1, 0.5));
      }

      if (ctx && gradient) ctx.fillStyle = gradient;
      if (ctx) ctx.fillRect(0, 0, size, size);

      if (ctx) ctx.fillStyle = "rgba(0, 0, 0, 255)";
      if (ctx) ctx.globalCompositeOperation = "destination-out";
      if (ctx) ctx.rotate(Math.PI / -4);
      for (let i = 0; i < size * 2; i += 32) {
        if (ctx) ctx.fillRect(-size, i, size * 2, 16);
      }

      return canvasDst;
    }

    // make 2 canvas elements, 300x300 with images in them.
    // We'll copy these to textures.
    const size = 600;
    const srcCanvas = createSourceImage(size);
    const dstCanvas = createDestinationImage(size);

    // Get a WebGPU context from the canvas and configure it
    const context = canvas.getContext("webgpu");
    const devicePixelRatio = window.devicePixelRatio;
    canvas.width = canvas.clientWidth * devicePixelRatio;
    canvas.height = canvas.clientHeight * devicePixelRatio;
    const presentationFormat = navigator.gpu.getPreferredCanvasFormat();

    const module = device.createShaderModule({
      label: "our hardcoded textured quad shaders",
      code: texturedQuadWGSL,
    });

    function createTextureFromSource(
      device: GPUDevice,
      source: HTMLCanvasElement,
      options: {
        flipY?: boolean;
        premultipliedAlpha?: boolean;
      } = {}
    ) {
      const { flipY, premultipliedAlpha } = options;
      const texture = device.createTexture({
        format: "rgba8unorm",
        size: [source.width, source.height],
        usage:
          GPUTextureUsage.TEXTURE_BINDING |
          GPUTextureUsage.COPY_DST |
          GPUTextureUsage.RENDER_ATTACHMENT,
      });
      device.queue.copyExternalImageToTexture(
        { source, flipY },
        { texture, premultipliedAlpha },
        { width: source.width, height: source.height }
      );
      return texture;
    }

    const bindGroupLayout = device.createBindGroupLayout({
      entries: [
        { binding: 0, visibility: GPUShaderStage.FRAGMENT, sampler: {} },
        { binding: 1, visibility: GPUShaderStage.FRAGMENT, texture: {} },
        { binding: 2, visibility: GPUShaderStage.VERTEX, buffer: {} },
      ],
    });

    const pipelineLayout = device.createPipelineLayout({
      bindGroupLayouts: [bindGroupLayout],
    });

    // create 2 textures with unpremultiplied alpha
    const srcTextureUnpremultipliedAlpha = createTextureFromSource(
      device,
      srcCanvas
    );
    const dstTextureUnpremultipliedAlpha = createTextureFromSource(
      device,
      dstCanvas
    );

    // create 2 textures with premultiplied alpha
    const srcTexturePremultipliedAlpha = createTextureFromSource(
      device,
      srcCanvas,
      { premultipliedAlpha: true }
    );
    const dstTexturePremultipliedAlpha = createTextureFromSource(
      device,
      dstCanvas,
      { premultipliedAlpha: true }
    );

    const sampler = device.createSampler({
      magFilter: "linear",
      minFilter: "linear",
      mipmapFilter: "linear",
    });

    type Uniforms = {
      buffer: GPUBuffer;
      values: Float32Array;
      matrix: Float32Array;
    };

    function makeUniformBufferAndValues(device: GPUDevice): Uniforms {
      // offsets to the various uniform values in float32 indices
      const kMatrixOffset = 0;

      // create a buffer for the uniform values
      const uniformBufferSize = 16 * 4; // matrix is 16 32bit floats (4bytes each)
      const buffer = device.createBuffer({
        label: "uniforms for quad",
        size: uniformBufferSize,
        usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
      });

      // create a typedarray to hold the values for the uniforms in JavaScript
      const values = new Float32Array(uniformBufferSize / 4);
      const matrix = values.subarray(kMatrixOffset, 16);
      return { buffer, values, matrix };
    }
    const srcUniform = makeUniformBufferAndValues(device);
    const dstUniform = makeUniformBufferAndValues(device);

    const srcBindGroupUnpremultipliedAlpha = device.createBindGroup({
      layout: bindGroupLayout,
      entries: [
        { binding: 0, resource: sampler },
        { binding: 1, resource: srcTextureUnpremultipliedAlpha.createView() },
        { binding: 2, resource: { buffer: srcUniform.buffer } },
      ],
    });

    const dstBindGroupUnpremultipliedAlpha = device.createBindGroup({
      layout: bindGroupLayout,
      entries: [
        { binding: 0, resource: sampler },
        { binding: 1, resource: dstTextureUnpremultipliedAlpha.createView() },
        { binding: 2, resource: { buffer: dstUniform.buffer } },
      ],
    });

    const srcBindGroupPremultipliedAlpha = device.createBindGroup({
      layout: bindGroupLayout,
      entries: [
        { binding: 0, resource: sampler },
        { binding: 1, resource: srcTexturePremultipliedAlpha.createView() },
        { binding: 2, resource: { buffer: srcUniform.buffer } },
      ],
    });

    const dstBindGroupPremultipliedAlpha = device.createBindGroup({
      layout: bindGroupLayout,
      entries: [
        { binding: 0, resource: sampler },
        { binding: 1, resource: dstTexturePremultipliedAlpha.createView() },
        { binding: 2, resource: { buffer: dstUniform.buffer } },
      ],
    });

    const textureSets = [
      {
        srcTexture: srcTexturePremultipliedAlpha,
        dstTexture: dstTexturePremultipliedAlpha,
        srcBindGroup: srcBindGroupPremultipliedAlpha,
        dstBindGroup: dstBindGroupPremultipliedAlpha,
      },
      {
        srcTexture: srcTextureUnpremultipliedAlpha,
        dstTexture: dstTextureUnpremultipliedAlpha,
        srcBindGroup: srcBindGroupUnpremultipliedAlpha,
        dstBindGroup: dstBindGroupUnpremultipliedAlpha,
      },
    ];

    const clearValue = [0, 0, 0, 0];
    const renderPassDescriptor: GPURenderPassDescriptor = {
      label: "our basic canvas renderPass",
      colorAttachments: [
        {
          view: undefined as unknown as GPUTextureView, // <- to be filled out when we render
          clearValue,
          loadOp: "clear",
          storeOp: "store",
        },
      ],
    };

    // Translates to/from a normalized color value and an 8bit unsigned color value.
    // This is because dat.gui only edits 8bit unsigned color values but we need normalized color values.

    function applyPreset() {
      const preset = presets[self.settings.preset];
      Object.assign(color, preset.color);
      Object.assign(alpha, preset.alpha || preset.color);
    }

    // const gui = new GUI({
    //   name: 'Blending',
    // });
    // this.gui = gui;
    // gui
    //   .add(settings, "alphaMode", ["opaque", "premultiplied"])
    //   .name("canvas alphaMode")
    //   .onChange(render);
    // gui
    //   .add(settings, "textureSet", [
    //     "premultiplied alpha",
    //     "un-premultiplied alpha",
    //   ])
    //   .name("texture data")
    //   .onChange(render);
    // gui.add(settings, "preset", [...Object.keys(presets)]).onChange(() => {
    //   applyPreset();
    //   render();
    // });

    // const colorFolder = gui.addFolder("color");
    // colorFolder.open();
    // colorFolder.add(color, "operation", operations).onChange(render);
    // colorFolder.add(color, "srcFactor", factors).onChange(render);
    // colorFolder.add(color, "dstFactor", factors).onChange(render);

    // const alphaFolder = gui.addFolder("alpha");
    // alphaFolder.open();
    // alphaFolder.add(alpha, "operation", operations).onChange(render);
    // alphaFolder.add(alpha, "srcFactor", factors).onChange(render);
    // alphaFolder.add(alpha, "dstFactor", factors).onChange(render);

    // const constantFolder = gui.addFolder("constant");
    // constantFolder.open();
    // constantFolder
    //   .addColor(new GUIColorHelper(constant.color), "value")
    //   .name("color")
    //   .onChange(render);
    // constantFolder.add(constant, "alpha", 0, 1).onChange(render);

    // const clearFolder = gui.addFolder("clear color");
    // clearFolder.open();
    // clearFolder.add(clear, "premultiply").onChange(render);
    // clearFolder.add(clear, "alpha", 0, 1).onChange(render);
    // clearFolder
    //   .addColor(new GUIColorHelper(clear.color), "value")
    //   .onChange(render);

    const dstPipeline = device.createRenderPipeline({
      label: "hardcoded textured quad pipeline",
      layout: pipelineLayout,
      vertex: {
        module,
      },
      fragment: {
        module,
        targets: [{ format: presentationFormat }],
      },
    });

    // @ts-ignore
    function makeBlendComponentValid(blend) {
      const { operation } = blend;
      if (operation === "min" || operation === "max") {
        blend.srcFactor = "one";
        blend.dstFactor = "one";
      }
    }

    function render() {
      makeBlendComponentValid(color);
      makeBlendComponentValid(alpha);
      // gui.updateDisplay();

      const srcPipeline = device?.createRenderPipeline({
        label: "hardcoded textured quad pipeline",
        layout: pipelineLayout,
        vertex: {
          module,
        },
        fragment: {
          module,
          targets: [
            {
              format: presentationFormat,
              blend: {
                color,
                alpha,
              },
            },
          ],
        },
      });

      const { srcTexture, dstTexture, srcBindGroup, dstBindGroup } =
        textureSets[self.settings.textureSet === "premultiplied alpha" ? 0 : 1];

      if (!device) {
        console.log("device not available");
        return;
      }
      context?.configure({
        device,
        format: presentationFormat,
        alphaMode: self.settings.alphaMode,
      });

      const canvasTexture = context?.getCurrentTexture();
      // Get the current texture from the canvas context and
      // set it as the texture to render to.
      // @ts-ignore
      renderPassDescriptor.colorAttachments[0].view =
        canvasTexture?.createView();

      // Apply the clearValue, pre-multiplying or not it based on the settings.
      {
        const { alpha, color, premultiply } = self.settings.clear;
        const mult = premultiply ? alpha : 1;
        clearValue[0] = color[0] * mult; // R
        clearValue[1] = color[1] * mult; // G
        clearValue[2] = color[2] * mult; // B
        clearValue[3] = alpha; // A
      }

      function updateUniforms(
        uniforms: Uniforms,
        canvasTexture: GPUTexture,
        texture: GPUTexture
      ) {
        const projectionMatrix = mat4.ortho(
          0,
          canvasTexture.width / devicePixelRatio,
          canvasTexture.height / devicePixelRatio,
          0,
          -1,
          1
        );

        mat4.scale(
          projectionMatrix,
          [texture.width, texture.height, 1],
          uniforms.matrix
        );

        // copy the values from JavaScript to the GPU
        device?.queue.writeBuffer(uniforms.buffer, 0, uniforms.values);
      }
      updateUniforms(srcUniform, canvasTexture!, srcTexture);
      updateUniforms(dstUniform, canvasTexture!, dstTexture);

      const encoder = device.createCommandEncoder({
        label: "render quad encoder",
      });
      const pass = encoder.beginRenderPass(renderPassDescriptor);

      // draw destination texture without blending
      pass.setPipeline(dstPipeline);
      pass.setBindGroup(0, dstBindGroup);
      pass.draw(6);

      // draw source texture with blending
      pass.setPipeline(srcPipeline!);
      pass.setBindGroup(0, srcBindGroup);
      pass.setBlendConstant([
        ...self.settings.constant.color,
        self.settings.constant.alpha,
      ]);
      pass.draw(6);

      pass.end();

      const commandBuffer = encoder.finish();
      device.queue.submit([commandBuffer]);
    }
    this.render = render;
    this.applyPreset = applyPreset;
    applyPreset();
    render();
  }

  onUpdate(data: any): void {
    // This sample does not require continuous updates.
    this.settings = data;
    this.settings.constant.color = new GUIColorHelper(
      typeof this.settings.constant.color === "string"
        ? set16ToRGBNormalized(this.settings.constant.color)
        : this.settings.constant.color
    ).value;
    this.settings.clear.color = new GUIColorHelper(
      typeof this.settings.clear.color === "string"
        ? set16ToRGBNormalized(this.settings.clear.color)
        : this.settings.clear.color
    ).value;
    if (this.applyPreset) {
      this.applyPreset();
    }
    // Re-render with the new settings.
    if (this.render) {
      this.render();
    }
  }

  destroy(): void {}
}
